/*
 *  ------------------------------------------------------------------
 *  Copyright © 2017 Hangzhou DtDream Technology Co.,Lt d. All rights reserved.
 *  ------------------------------------------------------------------
 *              Product: EMR
 *    Module Name: DataTrain
 *      Date Created: 2017-10-23
 *         Description:
 *  ------------------------------------------------------------------
 *  Modification History
 *  DATE            Name           Description
 *  ------------------------------------------------------------------
 *  2017-10-23     XZ 209
 *  ------------------------------------------------------------------
 */

package org.xukai.remoting.sdk.rpc.transfer;

import com.alibaba.fastjson.JSONObject;
import com.dtdream.emr.transmission.data.EmrSchema;
import com.dtdream.emr.transmission.exception.TransmissionException;
import com.dtdream.emr.transmission.io.Constant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.UUID;

import static com.dtdream.emr.transmission.exception.TransmissionException.Error.EXCEPTION_SDK_NETWORK;
import static com.dtdream.emr.transmission.exception.TransmissionException.Error.EXCEPTION_SDK_RETRY;
import static com.dtdream.emr.transmission.io.Constant.EMRSCHEMA;
import static com.dtdream.emr.transmission.io.Constant.JOBID;
import static com.dtdream.emr.transmission.io.Constant.PARTITION;
import static com.dtdream.emr.transmission.io.Constant.PROJECTNAME;
import static com.dtdream.emr.transmission.io.Constant.READERID;
import static com.dtdream.emr.transmission.io.Constant.TABLENAME;
import static com.dtdream.emr.transmission.io.Constant.URL_DOWNLOADJOB;
import static com.dtdream.emr.transmission.io.Constant.URL_JOB_CLOSE;
import static com.dtdream.emr.transmission.io.Constant.URL_READER;
import static com.dtdream.emr.transmission.io.Constant.URL_READER_CLOSE;
import static com.dtdream.emr.transmission.io.Constant.URL_READER_NUM;
import static com.dtdream.emr.transmission.io.Constant.URL_SCHEMA;
import static com.dtdream.emr.transmission.io.Constant.URL_UPLOADJOB;
import static com.dtdream.emr.transmission.io.Constant.URL_UPLOADJOB_COMMIT;
import static com.dtdream.emr.transmission.io.Constant.URL_WRITER;
import static com.dtdream.emr.transmission.io.Constant.URL_WRITER_CLOSE;
import static com.dtdream.emr.transmission.io.Constant.WRITERID;

public class TransferClient implements AutoCloseable {

    private static final Logger LOG = LoggerFactory.getLogger(TransferClient.class);

    private TransferNodePool pool;

    private String url;
    private NettyConfig config;
    private Set<TransferNode> nodeBlackList = new HashSet<>();
    private TransferNode currentNode;

    public TransferClient(String url, String accessKey, String keySecret, NettyConfig nettyConfig) {
        this.url=url;
        this.config=nettyConfig;
        this.pool = new TransferNodePool(accessKey, keySecret);
    }

    public void init() throws TransmissionException {
        pool.init(url, config);
        currentNode = pool.next();
    }

    private String convertString(String src) {
        String dst = src.replaceAll("=", "%3D");

        return dst;
    }

    private TransferNode assignNode() {
        return pool.next();
    }

    private Map<String, Object> postWithRetry(String url, Map<String, Object> params) throws TransmissionException {
        try {
            return currentNode.post(url, params);
        } catch (IOException e) {
            nodeBlackList.add(currentNode);
            if (nodeBlackList.size() == pool.size()) {
                LOG.error("Post failed on all DataTrainServers.");
                throw new TransmissionException(EXCEPTION_SDK_RETRY, e);
            }

            LOG.warn("Failed to post request to DataTrainServer:" + currentNode.address() + ". Try to post to another.");
            currentNode = pool.next();
            return postWithRetry(url, params);
        }
    }

    private TransferNode postWithLB(String url, Map<String, Object> params) throws TransmissionException {
        try {
            TransferNode node = assignNode();
            node.post(url, params);
            return node;
        } catch (IOException e) {
            nodeBlackList.add(currentNode);
            if (nodeBlackList.size() == pool.size()) {
                LOG.error("Post failed on all DataTrainServers.");
                throw new TransmissionException(EXCEPTION_SDK_RETRY, e);
            }

            LOG.warn("Failed to post request to DataTrainServer:" + currentNode.address() + ". Try to post to another.");
            return postWithLB(url, params);
        }
    }

    public EmrSchema getEmrSchema(UUID jobId) throws TransmissionException {
        Map<String, Object> params = new HashMap<>();

        params.put(JOBID, jobId);
        Map<String, Object> response = postWithRetry(URL_SCHEMA, params);
        JSONObject emrSchemaStr = (JSONObject) response.get(EMRSCHEMA);
        return JSONObject.parseObject(emrSchemaStr.toJSONString(), EmrSchema.class);
    }

    private UUID createJob(String path, String projectName, String tableName, String partition)
            throws TransmissionException {
        Map<String, Object> params = new HashMap<>();
        params.put(PROJECTNAME, projectName);
        params.put(TABLENAME, tableName);
        if (null != partition) {
            params.put(PARTITION, convertString(partition));
        }

        Map<String, Object> response = postWithRetry(path, params);
        String uuidStr = response.get(JOBID).toString();
        UUID uuid = UUID.fromString(uuidStr);

        return uuid;
    }

    public UUID createUploadJob(String projectName, String tableName, String partition) throws TransmissionException {
        return createJob(URL_UPLOADJOB, projectName, tableName, partition);
    }

    public UUID createUploadJob(String projectName, String tableName) throws TransmissionException {
        return createJob(URL_UPLOADJOB, projectName, tableName, null);
    }

    public void commitUploadJob(UUID jobId) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        postWithRetry(URL_UPLOADJOB_COMMIT, params);
    }

    public TransferNode createWriter(UUID jobId, String tableName, String patition, int writerId) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        params.put(WRITERID, writerId);
        return postWithLB(URL_WRITER, params);
    }

    public void closeWriter(UUID jobId, int writerId, TransferNode node) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        params.put(WRITERID, writerId);
        try {
            node.post(URL_WRITER_CLOSE, params);
        } catch (IOException e) {
            LOG.error("Close Writer failed", e);
            throw new TransmissionException(EXCEPTION_SDK_NETWORK);
        }
    }

    public UUID createDownloadJob(String projectName, String tableName, String partition) throws TransmissionException {
        return createJob(URL_DOWNLOADJOB, projectName, tableName, partition);
    }

    public UUID createDownloadJob(String projectName, String tableName) throws TransmissionException {
        return createJob(URL_DOWNLOADJOB, projectName, tableName, null);
    }

    public int getReaderNum(UUID jobId) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        Map<String, Object> response = postWithRetry(URL_READER_NUM, params);
        return Integer.valueOf(response.get(Constant.READERNUM).toString());
    }

    public TransferNode createReader(UUID jobId, int readerId) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        params.put(READERID, readerId);
        return postWithLB(URL_READER, params);
    }

    public void closeReader(UUID jobId, int readerId, TransferNode node) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        params.put(READERID, readerId);
        try {
            node.post(URL_READER_CLOSE, params);
        } catch (IOException e) {
            LOG.error("Close reader failed", e);
            throw new TransmissionException(EXCEPTION_SDK_NETWORK);
        }
    }

    public void closeJob(UUID jobId) throws TransmissionException {
        Map<String, Object> params = new HashMap<String, Object>();
        params.put(JOBID, jobId);
        postWithRetry(URL_JOB_CLOSE, params);
    }

    public void close() {
        pool.close();
    }
}
